#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @Time    : 2021/10/30 8:56 上午
# @Author  : WangZhixing

import json
from typing import Optional, Callable, List
import torch
from torch_geometric.data import Data, InMemoryDataset
from sklearn.preprocessing import LabelEncoder
from sklearn.preprocessing import OneHotEncoder
import random
from .DataInFo.DataInfo import DataInfo
from .Process import Process
from .Process.ProcessVector import ProcessVector


def index_to_mask(index, size):
    mask = torch.zeros((size,), dtype=torch.bool)
    mask[index] = 1
    return mask


class NolabelFileVectorDependenceGraph(InMemoryDataset):
    def __init__(self, root: str,
                 transform: Optional[Callable] = None,
                 pre_transform: Optional[Callable] = None):
        super().__init__(root, transform, pre_transform)
        self.data, self.slices = torch.load(self.processed_paths[3])
        with open(self.processed_paths[2], 'r') as f:
            self.dic = json.load(f)
        DataInfo(self.data)

    # 存放在root的目录下的raw文件夹中，如果没有调用下载
    @property
    def raw_file_names(self) -> List[str]:
        return ['edge.csv', 'new_filevectors.json']

    # 储存处理后的文件列表，若全部找到就跳过预处理。
    @property
    def processed_file_names(self) -> List[str]:
        return ['idx.edge.csv', 'idx.name.csv', "idx.name.json", 'data.pt']

    # 处理原始数据，并保存到processed_dir
    def process(self):
        process = ProcessVector()

        x = process.Node2Num(
            self.raw_paths[0],
            self.processed_paths[0],
            self.processed_paths[1],
            self.raw_paths[1]
        )

        # 节点的pairs，保存到edge_index
        edge_index = [[lst[0], lst[1]] for lst in process.lst]
        edge_index = torch.tensor(edge_index, dtype=torch.long)
        edge_index = edge_index - edge_index.min()
        edge_index = edge_index.t().contiguous()

        from torch_geometric.utils import to_undirected
        print("有向图边的数量", len(edge_index[0]))
        edge_index = to_undirected(edge_index)
        print("转化为无向图边的数量", len(edge_index[0]))
        # 保存边的权重
        # edge_attr = [lst[2] for lst in process.lst]
        # edge_attr = torch.tensor(edge_attr, dtype=torch.float)

        # 划分节点
        mask = [i for i in range(process.idx)]
        random.shuffle(mask)

        label = [lst for lst in process.NodeLabel.values()]
        label = torch.tensor(label, dtype=torch.long)

        idx_train = mask[0:int(len(mask) * 6 / 10)]
        idx_val = mask[int(len(mask) * 6 / 10):int(3 * len(mask) * 9 / 10)]
        idx_test = mask[int(3 * len(mask) / 10):]

        train_mask = index_to_mask(idx_train, len(mask))
        val_mask = index_to_mask(idx_val, len(mask))
        test_mask = index_to_mask(idx_test, len(mask))
        # 保存到data中
        data = Data(
            x=x
            , edge_index=edge_index
            , train_mask=train_mask,
            test_mask=test_mask, val_mask=val_mask
        )
        # 归一化
        data.x = data.x / data.x.sum(1, keepdim=True)

        # 数据，保存到pt文件中，方便下次读取
        torch.save(self.collate([data]), self.processed_paths[3])
        self.dic = process.dic
        with open(self.processed_paths[2], 'w') as f:
            json.dump(self.dic, f)
            # 打印data数据信息
